import os
import requests
import json
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Get API key from environment variables
API_KEY = os.getenv("DEEPSEEK_API_KEY")
if not API_KEY or API_KEY == "your_api_key_here":
    print("Please set your DeepSeek API key in the .env file")
    exit(1)

# DeepSeek API endpoint
API_URL = "https://api.deepseek.com/v1/chat/completions"

def generate_code(prompt, model="deepseek-coder", max_tokens=2000):
    """
    Generate code using the DeepSeek API
    
    Args:
        prompt (str): The coding task description
        model (str): The model to use (default: deepseek-coder)
        max_tokens (int): Maximum number of tokens to generate
        
    Returns:
        str: The generated code
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    # Create a system message to instruct the model to generate code
    messages = [
        {"role": "system", "content": "You are an expert programmer. Generate clean, efficient, and well-commented code based on the user's requirements."},
        {"role": "user", "content": prompt}
    ]
    
    data = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.2  # Lower temperature for more deterministic outputs
    }
    
    try:
        response = requests.post(API_URL, headers=headers, data=json.dumps(data))
        response.raise_for_status()
        
        result = response.json()
        if "choices" in result and len(result["choices"]) > 0:
            return result["choices"][0]["message"]["content"]
        else:
            return "Failed to generate code."
    except requests.exceptions.RequestException as e:
        return f"Error: {str(e)}"

def extract_code_blocks(text):
    """
    Extract code blocks from markdown text
    
    Args:
        text (str): Text containing markdown code blocks
        
    Returns:
        list: List of extracted code blocks
    """
    code_blocks = []
    lines = text.split('\n')
    in_code_block = False
    current_block = []
    language = ""
    
    for line in lines:
        if line.startswith('```') and not in_code_block:
            in_code_block = True
            # Extract language if specified
            language = line[3:].strip()
            continue
        elif line.startswith('```') and in_code_block:
            in_code_block = False
            code_blocks.append((language, '\n'.join(current_block)))
            current_block = []
            language = ""
            continue
        
        if in_code_block:
            current_block.append(line)
    
    return code_blocks

def save_code_to_file(code, filename):
    """
    Save code to a file
    
    Args:
        code (str): The code to save
        filename (str): The filename to save to
    """
    with open(filename, 'w') as f:
        f.write(code)
    print(f"Code saved to {filename}")

def main():
    print("DeepSeek Code Generation")
    print("Enter your coding task description, or type 'exit' to quit")
    print("-" * 50)
    
    while True:
        prompt = input("\nDescribe your coding task: ")
        
        if prompt.lower() == "exit":
            print("Goodbye!")
            break
        
        print("\nGenerating code...")
        generated_text = generate_code(prompt)
        
        print("\nGenerated Response:")
        print("-" * 50)
        print(generated_text)
        print("-" * 50)
        
        # Extract code blocks
        code_blocks = extract_code_blocks(generated_text)
        
        if code_blocks:
            print(f"\nFound {len(code_blocks)} code block(s).")
            for i, (language, code) in enumerate(code_blocks):
                print(f"\nCode Block {i+1} ({language}):")
                print("-" * 30)
                print(code)
                print("-" * 30)
                
                save_option = input(f"Save this code block? (y/n): ")
                if save_option.lower() == 'y':
                    filename = input("Enter filename to save to: ")
                    save_code_to_file(code, filename)
        else:
            print("\nNo code blocks found in the response.")

if __name__ == "__main__":
    main() 